#include "mainwindow.h"
#include "ui_mainwindow.h"
#include  <QImage>
#include  <QDebug>
#include  <QDateTime>
#include  <QThread>
#include  <QDir>
#include "v4l2.h"
#include "mpp.h"

extern "C"{
    #include <stdlib.h>
    #include <stdio.h>
    #include <string.h>
    #include <math.h>
    #include <time.h>
    #include <libavutil/avassert.h>
    #include <libavutil/channel_layout.h>
    #include <libavutil/opt.h>
    #include <libavutil/mathematics.h>
    #include <libavutil/timestamp.h>
    #include <libavformat/avformat.h>
    #include <libswscale/swscale.h>
    #include <libswresample/swresample.h>
    #include <stdio.h>
    #include <sys/ioctl.h>
    #include <linux/videodev2.h>
    #include <string.h>
    #include <sys/types.h>
    #include <sys/stat.h>
    #include <fcntl.h>
    #include <sys/mman.h>
    #include <poll.h>
    #include <stdlib.h>
    #include <stdio.h>
    #include <sys/types.h>
    #include <sys/socket.h>
    #include <arpa/inet.h>
    #include <pthread.h>
    #include <signal.h>
    #include <unistd.h>
    #include <string.h>
    #include <stdio.h>
}

//Yolo
static const unsigned char colors[19][3] = {
    {54, 67, 244}, {99, 30, 233}, {176, 39, 156}, {183, 58, 103}, {181, 81, 63},
    {243, 150, 33}, {244, 169, 3}, {212, 188, 0}, {136, 150, 0}, {80, 175, 76},
    {74, 195, 139}, {57, 220, 205}, {59, 235, 255}, {7, 193, 255}, {0, 152, 255},
    {34, 87, 255}, {72, 85, 121}, {158, 158, 158}, {139, 125, 96}
};

MainWindow::MainWindow(QWidget *parent)
    : QMainWindow(parent)
    , ui(new Ui::MainWindow)
    , timer(new QTimer(this))
    , isRecording(false)
    , isStreaming(false)
    , isYoloRunning(false)
    , yolo_initialized(false)
    
{
    ui->setupUi(this);

    QDir dir("/home/orangepi/saved_picture");
    if(!dir.exists()){
        dir.mkpath("/home/orangepi/saved_picture");
    }

    QDir videoDir("/home/orangepi/saved_video");
    if(!videoDir.exists()){
        videoDir.mkpath("/home/orangepi/saved_video");
    }
    if( v4l2_init(v4l2_ctx) == -1 ){
        qDebug() << "Initialization v4l2  failed";
        exit(-1);
    }

    //initial YOLOv11
    const char* model_path = "./yolo11.rknn"; // 替换为你的模型路径
    init_post_process();
    if (init_yolo11_model(model_path, &rknn_app_ctx) != 0) {
        qDebug() << "Failed to initialize YOLOv11 model:" << model_path;
    } else {
        yolo_initialized = true;
        qDebug() << "YOLOv11 model initialized successfully";
    }

    updateButtonStates();

    // Set up timer for frame capture (approx. 25 FPS)
    connect(timer, &QTimer::timeout, this, &MainWindow::captureAndDisplayFrame);
    timer->start(1000 / FPS);

    connect(ui->snapshotButton, &QPushButton::clicked, this, &MainWindow::snapshot);
    connect(ui->recordButton, &QPushButton::clicked, this, &MainWindow::startRecording);
    connect(ui->stopButton, &QPushButton::clicked, this, &MainWindow::stopRecording);
    connect(ui->startStreamButton, &QPushButton::clicked, this, &MainWindow::startStreaming);
    connect(ui->stopStreamButton, &QPushButton::clicked, this, &MainWindow::stopStreaming);
    connect(ui->startYoloButton, &QPushButton::clicked, this, &MainWindow::startYolo);
    connect(ui->stopYoloButton, &QPushButton::clicked, this, &MainWindow::stopYolo);
}

MainWindow::~MainWindow()
{
    timer->stop();
    stopRecording();
    stopStreaming();
    startYolo();
    if (yolo_initialized) {
        release_yolo11_model(&rknn_app_ctx);
        deinit_post_process();
        yolo_initialized = false;
    }
    v4l2_cleanup(v4l2_ctx);
    delete ui;
}

void MainWindow::updateButtonStates() {
    ui->recordButton->setEnabled(!isRecording);
    ui->stopButton->setEnabled(isRecording);
    ui->startStreamButton->setEnabled(!isStreaming);
    ui->stopStreamButton->setEnabled(isStreaming);
    ui->startYoloButton->setEnabled(!isYoloRunning);
    ui->stopYoloButton->setEnabled(isYoloRunning);
}

void MainWindow::captureAndDisplayFrame(){
    QMutexLocker locker(&mutex);
    
    struct v4l2_buffer buf;
    memset(&buf, 0, sizeof(buf));
    buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE_MPLANE;
    buf.memory = V4L2_MEMORY_MMAP;
    buf.m.planes = v4l2_ctx.planes.data(); //数组的大小
    buf.length = v4l2_ctx.planes.size();   //平面数组中的元素数

    // 取
    if (ioctl(v4l2_ctx.fd, VIDIOC_DQBUF, &buf) < 0) {
        qDebug() << "Failed to dequeue buffer:" << strerror(errno);
        return;
    }

    // Verify buffer data size
    if (buf.m.planes[0].bytesused != WIDTH * HEIGHT * 2) {
        qDebug() << "Warning: Buffer data size mismatch (" << buf.m.planes[0].bytesused << " bytes, expected " << WIDTH * HEIGHT * 2 << ")";
    }

    // Convert UYVY to RGB for display(qt)
    SwsContext *rgb_sws_ctx = sws_getContext(WIDTH, HEIGHT, AV_PIX_FMT_UYVY422,
                                             WIDTH, HEIGHT, AV_PIX_FMT_RGB24,
                                             SWS_BILINEAR, NULL, NULL, NULL);
    if (!rgb_sws_ctx) {
        qDebug() << "Failed to initialize RGB conversion context";
        return;
    }

    uint8_t *frame_data = static_cast<uint8_t*>(v4l2_ctx.buffers[buf.index].start);
    int src_linesize[1] = { WIDTH * 2 };
    uint8_t *dst_data[1] = { new uint8_t[WIDTH * HEIGHT * 3] };
    int dst_linesize[1] = { WIDTH * 3 };

    //UYVY422 --> RGB24
    sws_scale(rgb_sws_ctx, &frame_data, src_linesize, 0, HEIGHT, dst_data, dst_linesize);

    /*
    用于qt显示
    */
    // Create QImage 对象 from RGB(24) data 
    QImage image(dst_data[0], WIDTH, HEIGHT, QImage::Format_RGB888);//图像数据 宽度 高度 格式

    // 运行 YOLO 推理
    if (isYoloRunning && yolo_initialized) {
        processYoloInference(frame_data, image);
    }

    //将图像显示在 QLabel 中               将图像缩放到指定的大小  获取 QLabel 的大小 在缩放时保持图像的宽高比
    ui->label->setPixmap(QPixmap::fromImage(image.scaled(ui->label->size(), Qt::KeepAspectRatio)));

    /*
    用于MP4显示
    */
    // Encode and write frame to MP4 if recording
    if (isRecording) {
        encode_and_write_frame(frame_data, frameCount);
    }

    //stream 
    if (isStreaming) {
        encode_and_stream_frame(frame_data, frameCount);
    }

    // Clean up
    sws_freeContext(rgb_sws_ctx);
    delete[] dst_data[0];

    // Re-enqueue
    if (ioctl(v4l2_ctx.fd, VIDIOC_QBUF, &buf) < 0) {
        qDebug() << "Failed to re-enqueue buffer:" << strerror(errno);
        return;
    }
    frameCount++;
}

// yolo  绘制 yolo需要RGB888
void MainWindow::processYoloInference(uint8_t* frame_data, QImage& image)
{
    // 转换为 YOLO 模型输入（RGB888，调整尺寸）
    const int input_width = 320; // 假设模型输入尺寸
    const int input_height = 240;

    SwsContext* sws_ctx = sws_getContext(WIDTH, HEIGHT, AV_PIX_FMT_UYVY422,
        input_width, input_height, AV_PIX_FMT_RGB24,
        SWS_BILINEAR, NULL, NULL, NULL);
    
    if (!sws_ctx) {
        qDebug() << "Failed to initialize YOLO sws context";
        return;
    }

    uint8_t* rgb_data = new uint8_t[input_width * input_height * 3];
    int src_linesize[1] = { WIDTH * 2 };
    int dst_linesize[1] = { input_width * 3 };
    uint8_t* dst_data[1] = { rgb_data };
    sws_scale(sws_ctx, &frame_data, src_linesize, 0, HEIGHT, dst_data, dst_linesize);
    sws_freeContext(sws_ctx);

    //yolo输入
    image_buffer_t src_image;
    memset(&src_image, 0, sizeof(image_buffer_t));
    src_image.width  = input_width;
    src_image.height = input_height;
    src_image.format = IMAGE_FORMAT_RGB888;
    src_image.virt_addr = rgb_data;

    // rknn推理和处理
    object_detect_result_list od_results;
    if (inference_yolo11_model(&rknn_app_ctx, &src_image, &od_results) != 0) {
        qDebug() << "YOLOv11 inference failed";
        delete[] rgb_data;
        return;
    }

    //画框和概率
    QPainter painter(&image);
    int color_index = 0;
    for(int i = 0; i < od_results.count; i++){
        const unsigned char* color = colors[color_index % 19];
        color_index++;
        object_detect_result *det_result = &(od_results.results[i]);
        
        // 缩放边界框到原始图像尺寸
        float scale_x = (float)WIDTH / input_width;
        float scale_y = (float)HEIGHT / input_height;
        int left = det_result->box.left * scale_x;
        int top = det_result->box.top * scale_y;
        int right = det_result->box.right * scale_x;
        int bottom = det_result->box.bottom * scale_y;

        painter.setPen(QPen(QColor(color[0], color[1], color[2]), 2));
        painter.drawRect(left, top, right - left, bottom - top);

        QString text = QString("%1 %2%").arg(coco_cls_to_name(det_result->cls_id)).arg(det_result->prop * 100, 0, 'f', 1);
        QFont font("Arial", 10);
        painter.setFont(font);
        QRect textRect(left, top - 20, 100, 20);
        painter.fillRect(textRect, QColor(color[0], color[1], color[2]));
        painter.setPen(Qt::white);
        painter.drawText(textRect, Qt::AlignLeft | Qt::AlignVCenter, text);

        qDebug() << text << "@ (" << left << "," << top << "," << right << "," << bottom << ")";
    }
}

void MainWindow::snapshot(){
    QMutexLocker locker(&mutex);

    struct v4l2_buffer buf;
    memset(&buf, 0, sizeof(buf));
    buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE_MPLANE;
    buf.memory = V4L2_MEMORY_MMAP;
    buf.m.planes = v4l2_ctx.planes.data();
    buf.length = v4l2_ctx.planes.size();

    //取
    if (ioctl(v4l2_ctx.fd, VIDIOC_DQBUF, &buf) < 0) {
        qDebug() << "Failed to dequeue buffer for snapshot:" << strerror(errno);
        return;
    }

    QDir dir("/home/orangepi/saved_picture");
    QString filename = dir.filePath(QString("snapshot_%1.jpg").arg(QDateTime::currentDateTime().toString("yyyyMMdd_HHmmss")));
    //                                    内存映射的缓冲区
    if(save_snapshot(static_cast<uint8_t *>(v4l2_ctx.buffers[buf.index].start), filename.toStdString().c_str()) == 0){
        qDebug() << "Snapshot saved as" << filename;
    } else {
        qDebug() << "Failed to save snapshot";
    }

    // Re-enqueue buffer
    if (ioctl(v4l2_ctx.fd, VIDIOC_QBUF, &buf) < 0) {
        qDebug() << "Failed to re-enqueue buffer:" << strerror(errno);
    }
}

void MainWindow::startRecording(){
    QMutexLocker locker(&mutex);

    if(isRecording){
        qDebug() << "Already recording";
        return ;
    }

    QDir videoDir("/home/orangepi/saved_video");
    QString filename = videoDir.filePath(QString("video_%1.mp4")
                                        .arg(QDateTime::currentDateTime().toString("yyyyMMdd_HHmmss")));
    if(rkmpp_init(filename.toStdString().c_str()) == -1){
        qDebug() << "Failed to initialize FFmpeg";
        exit(-1);
    }

    isRecording = true;
    frameCount = 0;
    ui->recordButton->setEnabled(false);
    ui->stopButton->setEnabled(true);
    qDebug()<< "Started recording to" <<filename;
}

void MainWindow::stopRecording(){
    QMutexLocker locker(&mutex);

    if(!isRecording){
        qDebug() << "Not recording";
        return ;
    }

    isRecording = false;
    mpp_cleanup();
    frameCount = 0;
    ui->recordButton->setEnabled(true);
    ui->stopButton->setEnabled(false);
    qDebug()<< "Stop recording ";
}

void MainWindow::startStreaming(){
    QMutexLocker locker(&mutex);

    if(isStreaming){
        qDebug() << "已在推流中";
        return;
    }

    const char *output_url = "rtmp://47.110.53.247:1935/live/stream"; 
    if(rkmpp_init_stream(output_url) == -1){
        qDebug() << "无法初始化推流";
        return;
    }

    isStreaming = true;
    frameCount = 0;
    updateButtonStates();
    qDebug() << "开始推流到" << output_url;
}

void MainWindow::stopStreaming(){
    {
        QMutexLocker locker(&mutex);
        if(!isStreaming){
            qDebug() << "未在推流";
            return;
        }
        isStreaming = false;
        frameCount = 0;
    }

    mpp_cleanup_stream();
    updateButtonStates();
    qDebug() << "停止推流";
}

void MainWindow::startYolo()
{
    QMutexLocker locker(&mutex);
    if (isYoloRunning) {
        qDebug() << "YOLO already running";
        return;
    }
    if (!yolo_initialized) {
        qDebug() << "YOLO model not initialized";
        return;
    }
    isYoloRunning = true;
    updateButtonStates();
    qDebug() << "Started YOLO inference";
}

void MainWindow::stopYolo()
{
    QMutexLocker locker(&mutex);
    if (!isYoloRunning) {
        qDebug() << "YOLO not running";
        return;
    }
    isYoloRunning = false;
    updateButtonStates();
    qDebug() << "Stopped YOLO inference";
}

